//
// Created by wzc on 2023/2/6.
//
#include <ros/ros.h>
#include <geometry_msgs/PoseStamped.h>
#include <nav_msgs/Odometry.h>
#include <nav_msgs/Path.h>
#include <sensor_msgs/PointCloud2.h>
#include <pcl/point_cloud.h>
#include <pcl_conversions/pcl_conversions.h>
#include <math.h>


#include "plan_env/sdf_map.h"
#include "plan_env/map_ros.h"
#include "plan_env/edt_environment.h"

#include "path_searching/astar2.h"
#include "path_searching/kinodynamic_astar.h"

#include "plan_manage/planner_manager.h"
#include "plan_manage/plan_container.hpp"

#include "traj_utils/planning_visualization.h"

using namespace fast_planner;

Eigen::Vector3d start_pos, start_vel, start_acc;
Eigen::Vector3d end_pos, end_vel;

shared_ptr<FastPlannerManager> planner_manager_;
shared_ptr<PlanningVisualization> visualization_;

ros::Publisher kino_pub;
std::vector<Eigen::Vector3d> kino_path;

ros::Publisher astar_pub;
std::vector<Eigen::Vector3d> astar_path;

int drone_id = 1;
int drone_num = 3;

void init(ros::NodeHandle& nh){

  /* initialize main modules */
  planner_manager_.reset(new FastPlannerManager);
  planner_manager_->initPlanModules(nh);
  planner_manager_->swarm_traj_data_.init(drone_id, drone_num);

  visualization_.reset(new PlanningVisualization(nh));

  planner_manager_->edt_environment_->sdf_map_->setGlobalKnownFree();

  Eigen::Vector3d obs_min;
  Eigen::Vector3d obs_max;

  /// surrounding obs
  obs_min = Eigen::Vector3d(-20.0, -10.0, -0.5);
  obs_max = Eigen::Vector3d(20.0, -9.5, 3.0);
  planner_manager_->edt_environment_->sdf_map_->setObs(obs_min, obs_max);

  obs_min = Eigen::Vector3d(-20.0, 9.5, -0.5);
  obs_max = Eigen::Vector3d(20.0, 10.0, 3.0);
  planner_manager_->edt_environment_->sdf_map_->setObs(obs_min, obs_max);

  obs_min = Eigen::Vector3d(-20.0, -10.0, -0.5);
  obs_max = Eigen::Vector3d(-19.5, 10.0, 3.0);
  planner_manager_->edt_environment_->sdf_map_->setObs(obs_min, obs_max);

  obs_min = Eigen::Vector3d(19.5, -10.0, -0.5);
  obs_max = Eigen::Vector3d(20.0, 10.0, 3.0);
  planner_manager_->edt_environment_->sdf_map_->setObs(obs_min, obs_max);

  Eigen::Vector2d obs_center, obs_size;
  obs_size = Eigen::Vector2d(2.0, 2.0);

  obs_center = Eigen::Vector2d(9.0, 6.0);
  planner_manager_->edt_environment_->sdf_map_->setObsWithCenter(obs_center, obs_size);

  obs_center = Eigen::Vector2d(5.0, 0.0);
  planner_manager_->edt_environment_->sdf_map_->setObsWithCenter(obs_center, obs_size);

  obs_center = Eigen::Vector2d(13.0, -1.0);
  planner_manager_->edt_environment_->sdf_map_->setObsWithCenter(obs_center, obs_size);

  obs_center = Eigen::Vector2d(0.0, -4.0);
  planner_manager_->edt_environment_->sdf_map_->setObsWithCenter(obs_center, obs_size);

  obs_center = Eigen::Vector2d(-2.0, 3.0);
  planner_manager_->edt_environment_->sdf_map_->setObsWithCenter(obs_center, obs_size);

  obs_center = Eigen::Vector2d(-10.0, 5.0);
  planner_manager_->edt_environment_->sdf_map_->setObsWithCenter(obs_center, obs_size);

  obs_center = Eigen::Vector2d(-6.0, -1.0);
  planner_manager_->edt_environment_->sdf_map_->setObsWithCenter(obs_center, obs_size);

  obs_center = Eigen::Vector2d(-16.0, -0.0);
  planner_manager_->edt_environment_->sdf_map_->setObsWithCenter(obs_center, obs_size);

  planner_manager_->edt_environment_->sdf_map_->updateGlobalESDF3d();
  planner_manager_->edt_environment_->sdf_map_->setAllBoundToGlobal();

  start_pos << 0.0, 0.0, 0.6;
  start_vel << 0.5, 0.0, 0.0;
  start_acc << 0.0, 0.0, 0.0;

  end_vel << 0.0, 0.0, 0.0;
}

void testKinodynamicReplan(const Eigen::Vector3d& start_pos, const Eigen::Vector3d& start_vel, const Eigen::Vector3d& start_acc,
                           const Eigen::Vector3d& end_pos, const Eigen::Vector3d& end_vel) {
  /// =========== test kinodyanmicReplan =============== ///
  bool plan_success =
      planner_manager_->kinodynamicReplan(start_pos, start_vel, start_acc, end_pos, end_vel);

  /// =========== test planYaw =================== ///
  if (plan_success)
    planner_manager_->planYaw(Eigen::Vector3d(0.0, 0.0, 0.0));


  /// visualization
  // vis kino path
  nav_msgs::Path kino_path;
  kino_path.header.frame_id = "world";
  kino_path.poses.clear();
  for(auto & i : planner_manager_->plan_data_.kino_path_){
    geometry_msgs::PoseStamped pt;
    pt.pose.position.x = i(0);
    pt.pose.position.y = i(1);
    pt.pose.position.z = i(2);
    pt.pose.orientation.x = 0;
    pt.pose.orientation.y = 0;
    pt.pose.orientation.z = 0;
    pt.pose.orientation.w = 1;

    kino_path.poses.push_back(pt);
  }
  kino_pub.publish(kino_path);

  // 绘制当前无人机的轨迹
  auto local_data = &planner_manager_->local_data_;
  visualization_->drawBspline(local_data->position_traj_, 0.1,visualization_->getColor((drone_id-1)/double(drone_num)),
                              true, 0.25, Eigen::Vector4d(1, 1, 0, 1));
  auto plan_data = &planner_manager_->plan_data_;
  visualization_->drawYawTraj(local_data->position_traj_, local_data->yaw_traj_, plan_data->dt_yaw_);
}

void testPlanExploreTraj(const Eigen::Vector3d& start_pos, const Eigen::Vector3d& start_vel, const Eigen::Vector3d& start_acc,
                         const Eigen::Vector3d& end_pos, const Eigen::Vector3d& end_vel)  {
  planner_manager_->path_finder_->reset();
  if (planner_manager_->path_finder_->search(start_pos, end_pos, false) != Astar::REACH_END) {
    ROS_ERROR_STREAM("no path to next view");
    return;
  }

  astar_path = planner_manager_->path_finder_->getPath();
  std::cout << "astar path size " << astar_path.size() << std::endl;

  Eigen::Vector3d start_yaw = Eigen::Vector3d::Zero();
  double next_yaw = -2.0;
  double yd = 1.57;
  double diff0 = next_yaw - start_yaw[0];
  double diff1 = fabs(diff0);
  double time_lb = min(diff1, 2 * M_PI - diff1) / yd;

  ros::Time t1 = ros::Time::now();
  planner_manager_->planExploreTraj(astar_path, start_vel, start_acc, time_lb);

  std::cout << "traj duration " << planner_manager_->local_data_.position_traj_.duration_ << std::endl;

  double planTrajTime = (ros::Time::now() - t1).toSec();
  t1 = ros::Time::now();

  planner_manager_->planYawExplore(start_yaw, next_yaw, true, 1.0);


  /// visualization
  // vis astar path
  nav_msgs::Path ros_astar_path;
  ros_astar_path.header.frame_id = "world";
  ros_astar_path.poses.clear();
  for(auto & i : astar_path){
    geometry_msgs::PoseStamped pt;
    pt.pose.position.x = i(0);
    pt.pose.position.y = i(1);
    pt.pose.position.z = i(2);
    pt.pose.orientation.x = 0;
    pt.pose.orientation.y = 0;
    pt.pose.orientation.z = 0;
    pt.pose.orientation.w = 1;

    ros_astar_path.poses.push_back(pt);
  }
  astar_pub.publish(ros_astar_path);

  // 绘制当前无人机的轨迹
  auto local_data = &planner_manager_->local_data_;
  visualization_->drawBspline(local_data->position_traj_, 0.1,visualization_->getColor((drone_id-1)/double(drone_num)),
                              true, 0.25, Eigen::Vector4d(1, 1, 0, 1));
  auto plan_data = &planner_manager_->plan_data_;
  visualization_->drawYawTraj(local_data->position_traj_, local_data->yaw_traj_, plan_data->dt_yaw_);
}

void goalCallback(const geometry_msgs::PoseStampedConstPtr& goal){
  end_pos[0] = goal->pose.position.x;
  end_pos[1] = goal->pose.position.y;
  end_pos[2] = 1.8;

  std::cout << "==== start ====" << std::endl;
  std::cout << "p " << start_pos.transpose()  << " v " << start_vel.transpose() <<
  " a " << start_acc.transpose() << std::endl;
  std::cout << "==== end ====" << std::endl;
  std::cout << "p " << end_pos.transpose() << "v " << end_vel.transpose() << std::endl;

//  testKinodynamicReplan(start_pos, start_vel, start_acc, end_pos, end_vel);
  testPlanExploreTraj(start_pos, start_vel, start_acc, end_pos, end_vel);

  std::cout << "======== done =======" << std::endl;
}







int main(int argc, char** argv)
{
  ros::init(argc,argv,"test_plan_manage");
  ros::NodeHandle nh("~");

  init(nh);

  ros::Subscriber goal_sub = nh.subscribe("/move_base_simple/goal",10,goalCallback);

  kino_pub = nh.advertise<nav_msgs::Path>("/test/kino_path", 1000);
  astar_pub = nh.advertise<nav_msgs::Path>("/test/astar", 1000);



  ros::spin();
}